#include "pickdev.h"
#ifdef _WIN32
    #include <conio.h>
    #include <Windows.h>
    #include <winternl.h>
    #include <iostream>
    #include <boost/format.hpp>
    #include "utility.h"

    #if 0
    void DisplayVolumePaths(
            __in PWCHAR VolumeName
            )
    {
        DWORD  CharCount = MAX_PATH + 1;
        PWCHAR Names     = NULL;
        PWCHAR NameIdx   = NULL;
        BOOL   Success   = FALSE;

        for (;;)
        {
            //
            //  Allocate a buffer to hold the paths.
            Names = (PWCHAR) new BYTE [CharCount * sizeof(WCHAR)];

            if ( !Names )
            {
                //
                //  If memory can't be allocated, return.
                return;
            }

            //
            //  Obtain all of the paths
            //  for this volume.
            Success = GetVolumePathNamesForVolumeNameW(
                VolumeName, Names, CharCount, &CharCount
                );

            if ( Success )
            {
                break;
            }

            if ( GetLastError() != ERROR_MORE_DATA )
            {
                break;
            }

            //
            //  Try again with the
            //  new suggested size.
            delete [] Names;
            Names = NULL;
        }

        if ( Success )
        {
            //
            //  Display the various paths.
            for ( NameIdx = Names;
                  NameIdx[0] != L'\0';
                  NameIdx += wcslen(NameIdx) + 1 )
            {
                wprintf(L"  %s", NameIdx);
            }
            wprintf(L"\n");
        }

        if ( Names != NULL )
        {
            delete [] Names;
            Names = NULL;
        }

        return;
    }

    void print_partitions()
    {
        WCHAR volumeName[MAX_PATH];
        HANDLE findVolumeHandle = FindFirstVolumeW(volumeName, ARRAYSIZE(volumeName));
        verify(findVolumeHandle != INVALID_HANDLE_VALUE);
        while (true) {
            DWORD bytesReturned;
            struct {
                DWORD extentCount;
                DISK_EXTENT extents[64];
            } volumeDiskExtents;
            wprintf(L"%s\n", volumeName);
            volumeName[wcslen(volumeName) - 1] = L'\0';
            {
                HANDLE volumeHandle = CreateFileW(volumeName, 0, 0, NULL, OPEN_EXISTING, 0, NULL);
                verify(volumeHandle != INVALID_HANDLE_VALUE);
                if (!DeviceIoControl(
                        volumeHandle,
                        IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
                        NULL,
                        0,
                        &volumeDiskExtents,
                        sizeof(volumeDiskExtents),
                        &bytesReturned,
                        NULL)) {
                    DWORD error = GetLastError();
                    verify(CloseHandle(volumeHandle));
                    goto find_next;
                }
                verify(CloseHandle(volumeHandle));
            }
            assert(volumeDiskExtents.extentCount == 1);
            char diskName[MAX_PATH];
            sprintf(diskName, "\\\\?\\PhysicalDrive%u", volumeDiskExtents.extents[0].DiskNumber);
            HANDLE driveHandle = CreateFile(diskName, 0, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
            verify(INVALID_HANDLE_VALUE != driveHandle);
            char driveLayoutInfoBuf[sizeof(DRIVE_LAYOUT_INFORMATION_EX) + 15 * sizeof(PARTITION_INFORMATION_EX)];
            DRIVE_LAYOUT_INFORMATION_EX *const driveLayoutInfo = (DRIVE_LAYOUT_INFORMATION_EX *)driveLayoutInfoBuf;
            PARTITION_INFORMATION_EX (&partitionEntries)[16] = (PARTITION_INFORMATION_EX (&)[16])driveLayoutInfo->PartitionEntry;
            verify(DeviceIoControl(
                    driveHandle,
                    IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
                    NULL,
                    0,
                    driveLayoutInfoBuf,
                    sizeof(driveLayoutInfoBuf),
                    &bytesReturned,
                    NULL));
            verify(CloseHandle(driveHandle));
            for (size_t partitionIndex = 0; partitionIndex < driveLayoutInfo->PartitionCount; ++partitionIndex) {
                PARTITION_INFORMATION_EX const &pie(partitionEntries[partitionIndex]);
                printf("disk %u, partition %u, length %.1fGB, type %x\n",
                    volumeDiskExtents.extents[0].DiskNumber,
                    pie.PartitionNumber,
                    pie.PartitionLength.QuadPart / (double)(1024 * 1024 * 1024),
                    pie.Mbr.PartitionType);
            }
    find_next:
            if (!FindNextVolumeW(findVolumeHandle, volumeName, ARRAYSIZE(volumeName))) {
                assert(GetLastError() == ERROR_NO_MORE_FILES);
                break;
            }
        }

        verify(FindVolumeClose(findVolumeHandle));


    }
    #endif

    #if 0
    static void GetSizeString (LONGLONG size, wchar_t *str) {
        static wchar_t *b, *kb, *mb, *gb, *tb, *pb;
        struct {bool kilobyte; } opt = {
            true
        };
        if (b == NULL) {
            if (opt.kilobyte) {
                kb = L"KiB";
                mb = L"MiB";
                gb = L"GiB";
                tb = L"TiB";
                pb = L"PiB";
            } else {
                kb = L"KB";
                mb = L"MB";
                gb = L"GB";
                tb = L"TB";
                pb = L"PB";
            }
            b = L"bytes";
        }

        DWORD kilo = opt.kilobyte ? 1024 : 1000;
        LONGLONG kiloI64 = kilo;
        double kilod = kilo;

        if (size > kiloI64 * kilo * kilo * kilo * kilo * 99)
            swprintf (str, L"%I64d %s", size/ kilo / kilo /kilo/kilo/kilo, pb);
        else if (size > kiloI64*kilo*kilo*kilo*kilo)
            swprintf (str, L"%.1f %s",(double)(size/kilod/kilo/kilo/kilo/kilo), pb);
        else if (size > kiloI64*kilo*kilo*kilo*99)
            swprintf (str, L"%I64d %s",size/kilo/kilo/kilo/kilo, tb);
        else if (size > kiloI64*kilo*kilo*kilo)
            swprintf (str, L"%.1f %s",(double)(size/kilod/kilo/kilo/kilo), tb);
        else if (size > kiloI64*kilo*kilo*99)
            swprintf (str, L"%I64d %s",size/kilo/kilo/kilo, gb);
        else if (size > kiloI64*kilo*kilo)
            swprintf (str, L"%.1f %s",(double)(size/kilod/kilo/kilo), gb);
        else if (size > kiloI64*kilo*99)
            swprintf (str, L"%I64d %s", size/kilo/kilo, mb);
        else if (size > kiloI64*kilo)
            swprintf (str, L"%.1f %s",(double)(size/kilod/kilo), mb);
        else if (size > kiloI64)
            swprintf (str, L"%I64d %s", size/kilo, kb);
        else
            swprintf (str, L"%I64d %s", size, b);
    }


    static char *seconds_to_hhmmss(DWORD seconds, char *rv, int bufsiz) {
        DWORD hours = seconds / 3600;
        seconds -= hours * 3600;

        DWORD minutes = seconds / 60;
        seconds -= minutes * 60;

        _snprintf(rv, bufsiz, "%02d:%02d:%02d", hours, minutes, seconds);

        return rv;
    }
    // <adapted from truecrypt-4.2a-source-code/TrueCrypt/Common/Dlgcode.c>

    static int FakeDosNameForDevice (char *lpszDiskFile, char *lpszDosDevice, char *lpszCFDevice, BOOL bNameOnly) {
        if (strncmp(lpszDiskFile, "\\\\", 2) == 0) {
            strcpy(lpszCFDevice, lpszDiskFile);
            return 1;
        }

        BOOL bDosLinkCreated = TRUE;
        _snprintf(lpszDosDevice, MAX_PATH, "dskwipe%lu", GetCurrentProcessId());

        if (bNameOnly == FALSE)
            bDosLinkCreated = DefineDosDevice (DDD_RAW_TARGET_PATH, lpszDosDevice, lpszDiskFile);

        if (bDosLinkCreated == FALSE) {
            return 1;
        } else {
            _snprintf(lpszCFDevice, MAX_PATH, "\\\\.\\%s", lpszDosDevice);
        }

        return 0;
    }

    static int RemoveFakeDosName (char *lpszDiskFile, char *lpszDosDevice) {
        BOOL bDosLinkRemoved = DefineDosDevice (DDD_RAW_TARGET_PATH | DDD_EXACT_MATCH_ON_REMOVE |
                DDD_REMOVE_DEFINITION, lpszDosDevice, lpszDiskFile);
        if (bDosLinkRemoved == FALSE) {
            return 1;
        }

        return 0;
    }

    static void list_device(char *format_str, char *szTmp, int n) {
        int nDosLinkCreated;
        HANDLE dev;
        DWORD dwResult;
        BOOL bResult;
        PARTITION_INFORMATION diskInfo;
        DISK_GEOMETRY driveInfo;
        char szDosDevice[MAX_PATH], szCFDevice[MAX_PATH];
        static LONGLONG deviceSize = 0;
        wchar_t size[100] = {0}, partTypeStr[1024] = {0}, *partType = partTypeStr;

        BOOL drivePresent = FALSE;
        BOOL removable = FALSE;

        drivePresent = TRUE;

        //nDosLinkCreated = FakeDosNameForDevice (szTmp, szDosDevice,
        //	szCFDevice, FALSE);

        //dev = CreateFile (szCFDevice, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE , NULL, OPEN_EXISTING, FILE_ATTRIBUTE_READONLY, NULL);
        dev = CreateFile (szTmp, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE , NULL, OPEN_EXISTING, FILE_ATTRIBUTE_READONLY, NULL);

        bResult = DeviceIoControl (dev, IOCTL_DISK_GET_PARTITION_INFO, NULL, 0,
            &diskInfo, sizeof (diskInfo), &dwResult, NULL);

        // Test if device is removable
        if (/* n == 0 && */ DeviceIoControl (dev, IOCTL_DISK_GET_DRIVE_GEOMETRY, NULL, 0,
            &driveInfo, sizeof (driveInfo), &dwResult, NULL))
            removable = driveInfo.MediaType == RemovableMedia;

        RemoveFakeDosName(szTmp, szDosDevice);
        CloseHandle(dev);

        if (!bResult)
            return;

        // System creates a virtual partition1 for some storage devices without
        // partition table. We try to detect this case by comparing sizes of
        // partition0 and partition1. If they match, no partition of the device
        // is displayed to the user to avoid confusion. Drive letter assigned by
        // system to partition1 is displayed as subitem of partition0

        if (n == 0) {
            deviceSize = diskInfo.PartitionLength.QuadPart;
        }

        if (n > 0 && diskInfo.PartitionLength.QuadPart == deviceSize) {
            return;
        }

        switch(diskInfo.PartitionType) {
            case PARTITION_ENTRY_UNUSED:	partType = L""; break;
            case PARTITION_XINT13_EXTENDED:
            case PARTITION_EXTENDED:		partType = L"Extended"; break;
            case PARTITION_HUGE:			wsprintfW (partTypeStr, L"%s (0x%02X)", L"Unformatted", diskInfo.PartitionType); partType = partTypeStr; break;
            case PARTITION_FAT_12:			partType = L"FAT12"; break;
            case PARTITION_FAT_16:			partType = L"FAT16"; break;
            case PARTITION_FAT32:
            case PARTITION_FAT32_XINT13:	partType = L"FAT32"; break;
            case 0x08:						partType = L"DELL (spanning)"; break;
            case 0x12:						partType = L"Config/diagnostics"; break;
            case 0x11:
            case 0x14:
            case 0x16:
            case 0x1b:
            case 0x1c:
            case 0x1e:						partType = L"Hidden FAT"; break;
            case PARTITION_IFS:				partType = L"NTFS"; break;
            case 0x17:						partType = L"Hidden NTFS"; break;
            case 0x3c:						partType = L"PMagic recovery"; break;
            case 0x3d:						partType = L"Hidden NetWare"; break;
            case 0x41:						partType = L"Linux/MINIX"; break;
            case 0x42:						partType = L"SFS/LDM/Linux Swap"; break;
            case 0x51:
            case 0x64:
            case 0x65:
            case 0x66:
            case 0x67:
            case 0x68:
            case 0x69:						partType = L"Novell"; break;
            case 0x55:						partType = L"EZ-Drive"; break;
            case PARTITION_OS2BOOTMGR:		partType = L"OS/2 BM"; break;
            case PARTITION_XENIX_1:
            case PARTITION_XENIX_2:			partType = L"Xenix"; break;
            case PARTITION_UNIX:			partType = L"UNIX"; break;
            case 0x74:						partType = L"Scramdisk"; break;
            case 0x78:						partType = L"XOSL FS"; break;
            case 0x80:
            case 0x81:						partType = L"MINIX"; break;
            case 0x82:						partType = L"Linux Swap"; break;
            case 0x43:
            case 0x83:						partType = L"Linux"; break;
            case 0xc2:
            case 0x93:						partType = L"Hidden Linux"; break;
            case 0x86:
            case 0x87:						partType = L"NTFS volume set"; break;
            case 0x9f:						partType = L"BSD/OS"; break;
            case 0xa0:
            case 0xa1:						partType = L"Hibernation"; break;
            case 0xa5:						partType = L"BSD"; break;
            case 0xa8:						partType = L"Mac OS-X"; break;
            case 0xa9:						partType = L"NetBSD"; break;
            case 0xab:						partType = L"Mac OS-X Boot"; break;
            case 0xb8:						partType = L"BSDI BSD/386 swap"; break;
            case 0xc3:						partType = L"Hidden Linux swap"; break;
            case 0xfb:						partType = L"VMware"; break;
            case 0xfc:						partType = L"VMware swap"; break;
            case 0xfd:						partType = L"Linux RAID"; break;
            case 0xfe:						partType = L"WinNT hidden"; break;
            default:						wsprintfW(partTypeStr, L"0x%02X", diskInfo.PartitionType); partType = partTypeStr; break;
        }

        GetSizeString(diskInfo.PartitionLength.QuadPart, size);
        char *s_type = removable ? "Removable" : "Fixed";
        printf(format_str, szTmp, size, s_type, partType);
    }

    // </adapted from truecrypt-4.2a-source-code/TrueCrypt/Common/Dlgcode.c>

    void print_ticks(char *fmt, ULONGLONG ticks, ULONGLONG tick_frequency) {
        char wiping_time[255];
        DWORD seconds = (DWORD) (ticks / tick_frequency);
        seconds_to_hhmmss(seconds, wiping_time, sizeof(wiping_time));
        printf(fmt, wiping_time);
    }

    static void list_devices() {
        printf(
            "Device Name                         Size Type      Partition Type\n"
            "------------------------------ --------- --------- --------------------\n"
    //       123456789012345678901234567890 123456789 123456789 12345678901234567890
    //       \Device\Harddisk30\Partition03 1234.1 GB Removable SFS/LDM/Linux Swap
        );


        char *format_str = "%-30s %9S %-9s %-20S\n";

        char szTmp[MAX_PATH];
        int i;

        for (i = 0; i < 64; i++) {
            _snprintf(szTmp, sizeof(szTmp), "\\\\.\\PhysicalDrive%d", i);
            list_device(format_str, szTmp, 0);
        }

        for (i = 0; i < 64; i++) {
            for (int n = 0; n <= 32; n++) {
                _snprintf(szTmp, sizeof(szTmp), "\\Device\\Harddisk%d\\Partition%d", i, n);
                list_device(format_str, szTmp, n);
            }
        }

        for (i = 0; i < 8; i++) {
            _snprintf(szTmp, sizeof(szTmp), "\\Device\\Floppy%d", i);
            list_device(format_str, szTmp, 0);
        }

        list_device(format_str, "\\Device\\Ramdisk", 0);

        for (i = 0; i < 26; i++) {
            _snprintf(szTmp, sizeof(szTmp), "\\\\.\\%c:", 'A' + i);
            list_device(format_str, szTmp, 0);
        }
    }

    #endif

    bool choose_block_device(Device &device)
    {
        std::vector<Device const> devices;
        printf("    Drive Partition     Offset     Length Type\n");
        for (int driveIndex = 0; driveIndex < 64; ++driveIndex) {
            DWORD lastError = 0;
            char drivePath[MAX_PATH];
            sprintf_s(drivePath, "\\\\.\\PhysicalDrive%d", driveIndex);
            HANDLE driveHandle = CreateFile(drivePath, GENERIC_READ, 0, NULL, OPEN_EXISTING, 0, NULL);
            if (driveHandle == INVALID_HANDLE_VALUE) {
                lastError = GetLastError();
                continue;
            }
            DWORD bytesReturned;
            size_t const driveLayoutInfoSize = sizeof(DRIVE_LAYOUT_INFORMATION) + 63 * sizeof(PARTITION_INFORMATION);
            DRIVE_LAYOUT_INFORMATION *driveLayoutInfo = void_cast(
                DRIVE_LAYOUT_INFORMATION *,
                alloca(driveLayoutInfoSize));
            if (!DeviceIoControl(driveHandle, IOCTL_DISK_GET_DRIVE_LAYOUT, NULL, 0, driveLayoutInfo, driveLayoutInfoSize, &bytesReturned, NULL)) {
                lastError = GetLastError();
            } else {
                for (DWORD i = 0; i < driveLayoutInfo->PartitionCount; ++i) {
                    PARTITION_INFORMATION const &pi(driveLayoutInfo->PartitionEntry[i]);
                    if (pi.PartitionLength.QuadPart == 0)
                        continue;
                    switch (pi.PartitionType) {
                        case PARTITION_EXTENDED:
                        case PARTITION_FAT_16:
                        case PARTITION_FAT32:
                        case PARTITION_IFS: // ntfs?
                        case 0x83: // linux
                        case 0x82: // linux swap
                            continue;
                    }
                    Device const newDevice = {drivePath, pi.StartingOffset.QuadPart, pi.PartitionLength.QuadPart};
                    int32_t const MiB = 1024 * 1024;
                    std::cout << boost::format("%c) %5i %9i %7iMiB %7iMiB   %02x\n") % char(devices.size() + 'a') % driveIndex % i % (newDevice.offset / MiB) % (newDevice.size / MiB) % int(pi.PartitionType);
                    devices.push_back(newDevice);
                }
            }
            verify(CloseHandle(driveHandle));
        }
        if (devices.size()) {
            device = devices.at(getch() - 'a');
            return true;
        } else {
            printf("%s\n", "There are no devices to choose from");
            system("pause");
            return false;
        }
    }

    bool choose_file(Device &device)
    {
        OPENFILENAME ofn;
        memset(&ofn, 0, sizeof(ofn));
        ofn.lStructSize = sizeof(ofn);
        char path[MAX_PATH];
        path[0] = '\0';
        ofn.lpstrFile = path;
        ofn.nMaxFile = sizeof(path);
        if (GetOpenFileName(&ofn)) {
            device.path = path;
            device.offset = -1;
            device.size = -1;
            return true;
        } else {
            return false;
        }
    }


    bool choose_device(Device &device)
    {
        printf("%s",
            "Type of device?\n"
            "\t(f) file\n"
            "\t(b) block\n");
        verify(!fflush(stdin));
        switch (getch()) {
        case 'f':
            return choose_file(device);
        case 'b':
            return choose_block_device(device);
        default:
            return false;
        }
    }
#endif
